#include "tcp_conn.h"

#include "tcp_server.h"

namespace qio {

TcpConn::~TcpConn() QIO_NOEXCEPT {  // NOLINT
  delete fd_chan_;
  if (ssl_) {
    SSL_free(ssl_);
  }
}

TcpConn::TcpConn(EventLoop *loop, UniqueConfig &&config) QIO_NOEXCEPT
    : loop_(loop),
      conn_conf_(std::forward<UniqueConfig>(config)),
      fd_chan_(new FdChan(loop_, conn_conf_->fd)),
      conn_stat_(INIT),
      write_buf_queue_(),
      ssl_(nullptr),
      timeout_handle_(nullptr) {
  assert(loop_ != nullptr);
  assert(loop_->IsInLoop());
  assert(conn_conf_->read_buffer != nullptr &&
         conn_conf_->read_buffer_len != 0);
  // build ssl
  ConfigSSLContext();
}

void TcpConn::Start() QIO_NOEXCEPT {
  assert(loop_->IsInLoop());
  auto self = ShareInstance();
  assert(self);
  auto &conn_callbacks_ = conn_conf_->callbacks;
  const auto mode = conn_conf_->mode;
  if (mode == SERVER) {
    conn_stat_ = CONNECTED;
    conn_callbacks_.OnEstablished(this);
  } else {
    assert(mode == CLIENT);
    conn_stat_ = CONNECTING;
  }

  const auto disable_ssl = ssl_ == nullptr;
  fd_chan_->SetOnReadWriteCallback(std::bind(
      disable_ssl ? &TcpConn::HandleFdEvents : &TcpConn::HandleSSLFdEvents,
      self,  // point!
      std::placeholders::_1, std::placeholders::_2, std::placeholders::_3));

  fd_chan_->SetOnClosedCallback(std::bind(&TcpConn::HandleFdClosed,
                                          self,  // point!
                                          std::placeholders::_1));
  fd_chan_->SetOnErrorCallback(std::bind(
      &TcpConn::HandleFdError,
      self,  // point!
      std::placeholders::_1, std::placeholders::_2, std::placeholders::_3));

  fd_chan_->Start();
  fd_chan_->EnableReadWrite();
}

void TcpConn::Close() QIO_NOEXCEPT {
  if (loop_->IsInLoop()) {
    CloseInLoop();
  } else {
    loop_->RunInLoop(std::bind(&TcpConn::CloseInLoop, ShareInstance()));
  }
}

void TcpConn::Send(const SharedWriteBuf &buf) QIO_NOEXCEPT {
  if (loop_->IsInLoop()) {
    SendInLoop(buf);  // NOLINT
  } else {
    loop_->RunInLoop(std::bind(&TcpConn::SendInLoop, ShareInstance(), buf));
  }
}

void TcpConn::Send(const std::string &str) QIO_NOEXCEPT { Send(str, nullptr); }

void TcpConn::Send(const std::string &str,
                   WriteDoneCallback &&cb) QIO_NOEXCEPT {
  Send(str.c_str(), str.length(), std::forward<WriteDoneCallback>(cb));
}

void TcpConn::Send(const char *buf, std::size_t len) QIO_NOEXCEPT {
  Send(buf, len, nullptr);
}

void TcpConn::Send(const char *buf, std::size_t len,
                   WriteDoneCallback &&cb) QIO_NOEXCEPT {
  Send(std::make_shared<WriteBuf>(buf, len, false,
                                  std::forward<WriteDoneCallback>(cb)));
}

void TcpConn::SendInLoop(SharedWriteBuf buf) QIO_NOEXCEPT {
  if (conn_stat_ != CONNECTED) {
    return;
  }
  // 处于比较空闲的状态,可以尝试直接发送数据
  // 如果发送队列不为空,则表示 socket 处于一个 high watermark 状态
  // 可能无法一次性将数据发送出去, 因此将数据入队
  if (!write_buf_queue_.empty()) {
    write_buf_queue_.emplace_back(buf);
    fd_chan_->EnableWrite();
    return;
  }
  SendResult send_result = likely(ssl_ == nullptr) ? Write(buf) : SSLWrite(buf);
  switch (send_result) {
    case Ok: {
      return;
    }
    case Failed: {
      Close();
      return;
    }
    case Again: {
      write_buf_queue_.push_front(buf);
      fd_chan_->EnableWrite();
      return;
    }
    default: {
      assert(false);
    }  // unreachable code
  }
}

void TcpConn::CloseInLoop() QIO_NOEXCEPT {
  assert(loop_->IsInLoop());
  if (conn_stat_ >= CLOSING) {
    return;
  }
  conn_stat_ = CLOSING;
  if (ssl_) {
    // just mark ssl_ as shutdown
    // todo(suzhen): shutdown 后,可以选择接收对端回复的信号
    SSL_shutdown(ssl_);
  }
  fd_chan_->Close();
}

void TcpConn::HandleFdEvents(FdChan *ch, Fd fd, int events) QIO_NOEXCEPT {
  if (events & FdChan::WRITEABLE) {
    HandleWrite(ch, fd);
  }
  if (events & FdChan::READABLE) {
    ResetTimeoutHandle();
    HandleRead(ch, fd);
  }
}

void TcpConn::HandleSSLFdEvents(FdChan *ch, Fd fd, int events) QIO_NOEXCEPT {
  if (events & FdChan::WRITEABLE || SSL_want(ssl_) == SSL_READING) {
    HandleWrite(ch, fd);
  }
  if (events & FdChan::READABLE) {
    ResetTimeoutHandle();
    HandleSSLRead(ch, fd);
  }
}

TcpConn::SendResult TcpConn::SSLWrite(const SharedWriteBuf &buf) QIO_NOEXCEPT {
  const auto *data_will_send = buf->Data() + buf->WriteIndex();
  const auto data_will_send_len =
      static_cast<int>(buf->DataLen() - buf->WriteIndex());
  const auto send_len = SSL_write(ssl_, data_will_send, data_will_send_len);
  if (send_len == data_will_send_len) {
    buf->SetWriteIndex(static_cast<std::size_t>(data_will_send_len));
    return Ok;
  }
  if (send_len <= 0) {
    switch (SSL_get_error(ssl_, send_len)) {
      case SSL_ERROR_WANT_READ:
      case SSL_ERROR_WANT_WRITE:
        return Again;
      default: {
        return Failed;
      }  // ssl fatal error
    }
  }
  // partial send
  buf->SetWriteIndex(buf->WriteIndex() + send_len);
  return Again;
}

TcpConn::SendResult TcpConn::Write(const SharedWriteBuf &buf) QIO_NOEXCEPT {
  const auto *data_will_send = buf->Data() + buf->WriteIndex();
  const auto data_will_send_len = buf->DataLen() - buf->WriteIndex();
  // write direct
  const auto send_len = socket_send(fd_chan_->GetFd(), data_will_send,
                                    data_will_send_len, MSG_NOSIGNAL);
  if (send_len == static_cast<ssize_t>(data_will_send_len)) {
    // send ok
    buf->SetWriteIndex(data_will_send_len);
    return Ok;
  }
  if (send_len < 0) {
    if (IO_ERR_RW_RETRIABLE(GetLastSocketError())) {
      // enqueue data
      return Again;
    }
    // socket error
    return Failed;
  }
  // partial send
  buf->SetWriteIndex(buf->WriteIndex() + send_len);
  return Again;
}

void TcpConn::ConfigSSLContext() QIO_NOEXCEPT {
  auto ssl_helper = conn_conf_->ssl_helper;
  const auto mode = conn_conf_->mode;
  if (!ssl_helper) {
    return;
  }
  auto *ssl_ctx = ssl_helper->GetSSLContext();
  if (ssl_ctx != nullptr) {
    ssl_ = SSL_new(ssl_ctx);
  }
  if (ssl_ != nullptr) {
    assert(ssl_ != nullptr);
    SSL_set_fd(ssl_, fd_chan_->GetFd());
    SSL_set_mode(ssl_, SSL_MODE_RELEASE_BUFFERS);
    if (mode == SERVER) {
      // server端设置 accept_state
      SSL_set_accept_state(ssl_);
    } else {
      // client端设置 connect_state
      SSL_set_connect_state(ssl_);
    }
  }
}

void TcpConn::HandleRead(FdChan *, Fd fd) QIO_NOEXCEPT {
  if (conn_stat_ != CONNECTED) {
    return;
  }
  auto *read_buf_ = conn_conf_->read_cache_buffer.get();
  auto *read_buffer_ = conn_conf_->read_buffer;
  auto read_buffer_len_ = conn_conf_->read_buffer_len;
  auto &conn_callbacks_ = conn_conf_->callbacks;
  const auto recv_len = socket_recv(fd, read_buffer_, read_buffer_len_, 0);
  if (recv_len > 0) {
    read_buf_->Write(read_buffer_, static_cast<std::size_t>(recv_len));
    conn_callbacks_.OnMessage(this, read_buf_);
  } else if (recv_len == 0) {
    // EOF special
    Close();
  } else if (recv_len == -1) {
    if (IO_ERR_RW_RETRIABLE(GetLastSocketError())) {
      fd_chan_->EnableRead();
    } else {
      Close();
      return;
    }
  } else {
    // unreachable code
    assert(false);
  }
}

void TcpConn::HandleSSLRead(FdChan *, Fd) QIO_NOEXCEPT {
  auto *read_buf_ = conn_conf_->read_cache_buffer.get();
  auto *read_buffer_ = conn_conf_->read_buffer;
  auto read_buffer_len_ = conn_conf_->read_buffer_len;
  auto &conn_callbacks_ = conn_conf_->callbacks;

  do {
    if (conn_stat_ != CONNECTED) {
      return;
    }
    const int read_result =
        SSL_read(ssl_, read_buffer_, static_cast<int>(read_buffer_len_));
    if (read_result <= 0) {
      switch (SSL_get_error(ssl_, read_result)) {
        case SSL_ERROR_WANT_READ: {
        } break;
        case SSL_ERROR_WANT_WRITE: {
          fd_chan_->EnableWrite();
        } break;
        default: {
          Close();
          return;
        }
      }
      break;
    } else {
      read_buf_->Write(read_buffer_, static_cast<std::size_t>(read_result));
      conn_callbacks_.OnMessage(this, read_buf_);
    }
  } while (SSL_pending(ssl_));
}

void TcpConn::HandleWrite(FdChan *, Fd) QIO_NOEXCEPT {
  if (conn_stat_ != CONNECTED) {
    return;
  }
  if (write_buf_queue_.empty()) {
    fd_chan_->DisableWrite();
    return;
  }
  const auto &buf = write_buf_queue_.front();
  const auto send_result = likely(ssl_ == nullptr) ? Write(buf) : SSLWrite(buf);
  switch (send_result) {
    case Ok: {
      write_buf_queue_.pop_front();
      if (write_buf_queue_.empty()) {
        fd_chan_->DisableWrite();
      }
    } break;
    case Again: {
      fd_chan_->EnableWrite();
      return;
    }
    case Failed: {
      // notify send fail
      write_buf_queue_.pop_front();
      Close();
      return;
    }
    default: {
      assert(false);
    }
  }
}

void TcpConn::HandleFdClosed(FdChan *) QIO_NOEXCEPT {
  auto &conn_callbacks_ = conn_conf_->callbacks;
  conn_stat_ = CLOSED;
  conn_callbacks_.OnDisconnected(this);
  conn_callbacks_.Reset();
}

void TcpConn::HandleFdError(FdChan *, int err_code,
                            const char *err_msg) QIO_NOEXCEPT {
  auto &conn_callbacks_ = conn_conf_->callbacks;
  conn_callbacks_.OnError(this, err_code, err_msg);
}

}  // namespace qio
